import openai
import faiss
import numpy as np
import sqlite3
from typing import List, Tuple

# 初始化OpenAI客户端
client = openai.OpenAI(
    api_key="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"
)


class RAGRetriever:
    def __init__(self, db_path: str = "openai_rag_database.db", index_path: str = "openai_rag_index.faiss"):
        """
        初始化RAG检索系统
        :param db_path: SQLite数据库路径
        :param index_path: FAISS索引路径
        """
        # 加载FAISS向量索引
        self.index = faiss.read_index(index_path)

        # 连接SQLite数据库
        self.conn = sqlite3.connect(db_path)
        self.cursor = self.conn.cursor()

        # 验证数据库是否有数据
        self.cursor.execute("SELECT COUNT(*) FROM documents")
        count = self.cursor.fetchone()[0]
        if count == 0:
            raise ValueError("数据库中没有文档，请先添加文档")
        print(f"数据库已加载，共有 {count} 个文档")

        # 获取向量维度
        self.dimension = self.index.d

    def get_embedding(self, text: str, model: str = "text-embedding-3-small") -> np.ndarray:
        """
        使用OpenAI API获取单个文本的嵌入向量
        """
        response = client.embeddings.create(input=[text], model=model)
        return np.array(response.data[0].embedding)

    def retrieve(self, query: str, k: int = 3) -> List[Tuple[str, float]]:
        """
        检索与查询最相关的k个文档
        :param query: 查询文本
        :param k: 返回结果数量
        :return: 包含(文档文本, 相似度分数)的列表
        """
        # 获取查询向量
        query_embedding = self.get_embedding(query)
        query_embedding = query_embedding.reshape(1, -1)  # 转换为2D数组

        # 在FAISS中搜索
        distances, indices = self.index.search(query_embedding, k)

        # 从数据库获取原始文本
        results = []
        for idx, dist in zip(indices[0], distances[0]):
            # FAISS返回的是0-based索引，我们数据库id是1-based
            self.cursor.execute(f"SELECT text FROM documents WHERE id={idx + 1}")
            result = self.cursor.fetchone()
            if result is None:
                print(f"警告: 未找到ID {idx + 1} 的文档")
                continue
            text = result[0]
            similarity_score = 1 / (1 + dist)  # 将距离转换为相似度分数(0-1)
            results.append((text, similarity_score))

        # 按相似度降序排序
        results.sort(key=lambda x: x[1], reverse=True)
        return results

    def close(self):
        """关闭数据库连接"""
        self.conn.close()


# 使用示例
if __name__ == "__main__":
    retriever = RAGRetriever()

    # 示例查询
    queries = ["Christmas Martini是什么?", "Christmas Martini怎么制作?", "我想喝点甜的酒"]

    for query in queries:
        print(f"\n查询: '{query}'")
        results = retriever.retrieve(query, k=2)

        for i, (text, score) in enumerate(results, 1):
            print(f"\n结果 {i} (相似度: {score:.2f}):")
            print(text)

    retriever.close()
